import numpy as np  
import matplotlib.pyplot as plt  
import math
from filterpy.kalman import ExtendedKalmanFilter as EKF  
from scipy.spatial.transform import Rotation as R  
  
# 扩展卡尔曼滤波器模型定义  
def f(x, dt):  
    """State transition function for nearly constant velocity."""  
    return np.array([x[0]+x[3]*dt+0.5*x[6]*dt*dt,x[1]+x[4]*dt+0.5*x[7]*dt*dt,x[2]+x[5]*dt+0.5*x[8]*dt*dt, x[3]+x[6]*dt, x[4]+x[7]*dt, x[5]+x[8]*dt,x[6],x[7],x[8]])   
  
def h(x):  
    # 直接观测位置  
    return x[:3]  
  
 
def H(x):  
     
    return np.array([[1, 0, 0, 0, 0, 0, 0, 0, 0],  
                     [0, 1, 0, 0, 0, 0, 0, 0, 0],  
                     [0, 0, 1, 0, 0, 0, 0, 0, 0]])  

# 读取txt文件x y z  
def read_data(filename):  
    with open(filename, 'r') as file:  
        lines = file.readlines()  
        data = []  
        for line in lines:  
            # 假设每行数据都是三个以空格分隔的数字  
            x, y, z = line.strip().split()  # 使用split()函数分割字符串，这里假设数据以空格分隔  
            data.append([float(x), float(y), float(z)])  # 使用float()函数将字符串转换为浮点数  
    return np.array(data)   
# 读取txt文件vx vy vz  
def read_data_v(filename):  
    with open(filename, 'r') as file:  
        lines = file.readlines()  
        data = []  
        data_v = []
        x_old = -3.58035  
        y_old = 15.4415
        z_old = 21.7189
        vx_old = 0
        vy_old = 0
        vz_old =0
        ax_old = 0
        ay_old = 0
        az_old =0
        for line in lines:  
            # 假设每行数据都是三个以空格分隔的数字  
            x, y, z = line.strip().split()  # 使用split()函数分割字符串，这里假设数据以空格分隔  
            x_new = float(x)
            y_new = float(y)
            z_new = float(z)
            vx = (x_new-x_old)/0.1
            vy = (y_new-y_old)/0.1
            vz = (z_new-z_old)/0.1  # 使用float()函数将字符串转换为浮点数 
            ax = (vx - vx_old)/0.1
            ay = (vy - vy_old)/0.1
            az = (vz - vz_old)/0.1
            data.append([float(x), float(y), float(z)])  # 使用float()函数将字符串转换为浮点数  
            data_v.append([vx, vy, vz])  # 使用float()函数将字符串转换为浮点数
            x_old = x_new
            y_old = y_new
            z_old = z_new
            vx_old = vx
            vy_old = vy
            vz_old = vz
            ax_old = ax
            ay_old = ay
            az_old = az
    return np.array(data)   
  
# 绘制3D散点图  
def plot_3d(data, title="3D Scatter Plot"):  
    fig = plt.figure()  
    ax = fig.add_subplot(111, projection='3d')  
    #ax.scatter(data[:,0], data[:,1], data[:,2])  
    ax.plot(data[:,0], data[:,1], data[:,2])  
    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('Z')  
    ax.set_title(title)  
    plt.show()  

  
if __name__ == "__main__":  
    filename = "dynamic_target_coord.txt"  # 替换为你的txt文件名  
    #filename = "1.txt"
    #data = read_data_v(filename)  
    #plot_2d(data)
    data_txt = read_data_v(filename)
    # 设置时间、频率和振幅参数  
    t_total = 34  # 总时间（秒）  
    frequency = 10  # 频率（Hz）  
    dt = 1 / frequency  # 时间步长（秒）  
    num_points = 340  # 总点数  
    times = np.arange(0, t_total, dt)  # 时间数组  
    qconstant = 0.1
    rconstant = 1.2

    #data


    ekf = EKF(dim_x=9, dim_z=3)  # 状态向量为[x, y, z]，测量向量为[x, y, z]  
    ekf.x = np.array([data_txt[0,0], data_txt[0,1], data_txt[0,2], 0, 0, 0,0,0,0])  # 初始状态（位置和速度）设置为第一个测量值，速度为0  
    ekf.P = np.eye(9) * 5**2  # 初始状态协方差矩阵，对角线上的值表示对应状态的不确定性程度（这里假设所有状态的不确定性相同）  

    ekf.R = np.eye(3) * (rconstant)   # 测量噪声协方差矩阵，设置为噪声的标准差的平方（转换为厘米） 
    ekf.Q = np.array([[10*qconstant, 0, 0, 0, 0, 0, 0, 0, 0],  
                      [0, 10*qconstant, 0, 0, 0, 0, 0, 0, 0],  
                      [0, 0, 10*qconstant, 0, 0, 0, 0, 0, 0],  
                      [0, 0, 0, qconstant, 0, 0, 0, 0, 0],  
                      [0, 0, 0, 0, qconstant, 0, 0, 0, 0],  
                      [0, 0, 0, 0, 0, qconstant, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0.1*qconstant, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0.1*qconstant, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0.1*qconstant]]) # 过程噪声协方差矩阵，这里假设过程噪声较小，因此设置为较小的值（可以根据实际情况调整）  
    ekf.F = np.array([[1, 0, 0, dt, 0, 0, 0.5 * dt**2 , 0, 0],  
                      [0, 1, 0, 0, dt, 0, 0, 0.5 * dt**2 , 0],  
                      [0, 0, 1, 0, 0, dt, 0, 0, 0.5 * dt**2 ],  
                      [0, 0, 0, 1, 0, 0, dt, 0, 0],  
                      [0, 0, 0, 0, 1, 0, 0, dt, 0],  
                      [0, 0, 0, 0, 0, 1, 0, 0, dt],
                      [0, 0, 0, 0, 0, 0, 1, 0, 0],  
                      [0, 0, 0, 0, 0, 0, 0, 1, 0],  
                      [0, 0, 0, 0, 0, 0, 0, 0, 1]
                      ]) # 状态转移函数的雅可比矩阵（使用lambda函数传递dt）  
    ekf.H = H  # 测量函数的雅可比矩阵（不需要状态向量，因此直接传递函数）  
    ekf.f = f  # 状态转移函数（使用lambda函数传递dt）  
    ekf.h = h  # 测量函数（不需要状态向量，因此直接传递函数） 
    filtered_data_nom = []  # 用于存储滤波后的数据  

    for tt in data_txt:  # 对每个测量值进行滤波处理  

        ekf.predict()  # 使用状态转移函数预测下一个时间步的状态估计和协方差矩阵（尽管这里我们实际上是在处理当前时间步的下一个测量值）
        print("predict")
        #print(ekf.x[:3])
        ekf.update(tt,HJacobian=ekf.H,Hx=ekf.h)
        print("update")
        #print(ekf.x[:3])
        print("---------->")
          # 使用当前测量值更新滤波器状态估计和协方差矩阵  
        filtered_data_nom.append(ekf.x)  # 存储滤波后的位置估计（忽略速度估计）  
    filtered_data_nom = np.array(filtered_data_nom)  # 将列表转换为NumPy数组以便于绘图处理  
    fig = plt.figure()  

    # Cam
    ax1 = fig.add_subplot(121, projection='3d')  
    ax1.plot(filtered_data_nom[:, 0], filtered_data_nom[:, 1], filtered_data_nom[:, 2],  label='filtered_data_nom')  # 原图为蓝色  

    ax1.set_xlabel("X Position (cm)")  
    ax1.set_ylabel("Y Position (cm)")  
    ax1.set_zlabel("Z Position (cm)")  
    ax1.set_title("Trajectory after ekf")  
    plt.legend()  
    ax1 = fig.add_subplot(122, projection='3d')  
    ax1.plot(data_txt[:, 0], data_txt[:, 1], data_txt[:, 2], 'pink',label='Original')  # 原图为蓝色  

    ax1.set_xlabel("X Position (cm)")  
    ax1.set_ylabel("Y Position (cm)")  
    ax1.set_zlabel("Z Position (cm)")  
    ax1.set_title("Trajectory ")  
    plt.legend()  


    plt.show()


